package deepbrain.iit;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import deepbrain.simnetwork.network.Network;
import deepbrain.simnetwork.network.NetworkState;
import deepbrain.simnetwork.network.NodeState;
import deepbrain.simnetwork.util.SimNetUtils;

/**
 * Compute {PHI} for a given network. This is only prototype, which do not
 * consist any optimization of algorithms nor caching mechanisms.
 * 
 * @author Li Yang
 * @create 2009-3-1
 */
public class PhiComputation {

	private final Network network;

	private int numOfNodes;
	// private int numOfStates;

	private List<Integer> totalPart = new ArrayList<Integer>();

	private Map<NetworkState, HashMap<NetworkState, Double>> p_X0_X1 = new HashMap<NetworkState, HashMap<NetworkState, Double>>();

	private Map<NetworkState, Double> systemEIs = new HashMap<NetworkState, Double>();

	public Network getNetwork() {
		return network;
	}

	public PhiComputation(Network network) {
		this.network = network;
		this.numOfNodes = network.getNumOfNodes();
		// this.numOfStates = 1 << numOfNodes;

		for (int i = 0; i < numOfNodes; i++)
			totalPart.add(i);
	}

	public PhiCompResult findMIP(NetworkState state,
			Iterable<Partition<Integer>> partitions) {

 		if (!p_X0_X1.containsKey(state)) {
			p_X0_X1.put(state, computeP_X0_X1(state));
			systemEIs.put(state, compSystemEI(state));
		}

		PhiCompResult result = new PhiCompResult();
		double min_phi = Double.MAX_VALUE;
		double norm_ei;
		EICompResult eiResult;
		for (Partition<Integer> partition : partitions) {
			eiResult = compEI4Partition(state, partition);
			norm_ei = eiResult.ei / eiResult.norm;
//			System.out.println(eiResult);
// Display all arbitrary PHI
			result.MIPs.add(eiResult);
			if (norm_ei < min_phi) {
				result.MIN_NORMED_EI = norm_ei;
				result.PHI = eiResult.ei;
				min_phi = norm_ei;
			} else if (norm_ei == min_phi) {
				if (eiResult.ei < result.PHI)
					result.PHI = eiResult.ei;
			}
			
/*			
			if (norm_ei < min_phi) {
				result.MIPs.clear();
				result.MIPs.add(eiResult);
				result.MIN_NORMED_EI = norm_ei;
				result.PHI = eiResult.ei;
				min_phi = norm_ei;
			} else if (norm_ei == min_phi) {
				result.MIPs.add(eiResult);
				if (eiResult.ei < result.PHI)
					result.PHI = eiResult.ei;
			}
*/
		}
		return result;
	}

	/**
	 * Find MIPs in bipartitions
	 */
	public PhiCompResult findMIP(NetworkState state) {
		PartitionGenerator partitions = new PartitionGenerator(numOfNodes, 2);
		return findMIP(state, partitions);
	}

	/**
	 * Find MIPs in given max num of partitions
	 */
	public PhiCompResult findMIP(NetworkState state, int maxNumOfParts) {
		PartitionGenerator partitions = new PartitionGenerator(numOfNodes,
				maxNumOfParts);
		return findMIP(state, partitions);
	}

	private EICompResult compEI4Partition(NetworkState state,
			Partition<Integer> partition) {
		double norm = normalization(partition);
		if (partition.numOfParts() == 1)
			return new EICompResult(systemEIs.get(state), norm, partition);
		List<HashMap<NetworkState, Double>> pM0M1s = compM0M1s(state, partition);
		HashMap<NetworkState, Double> pX0X1 = p_X0_X1.get(state);
		return new EICompResult(computeEI(pX0X1, pM0M1s), norm, partition);
	}

	private final List<HashMap<NetworkState, Double>> compM0M1s(
			NetworkState networkState, Partition<Integer> partition) {
		List<HashMap<NetworkState, Double>> pM0M1s = new ArrayList<HashMap<NetworkState, Double>>();
		for (int i = 0; i < partition.numOfParts(); i++) {
			pM0M1s.add(computeP_M0_M1(networkState, partition.part(i)));
		}
		return pM0M1s;
	}

	private final double compSystemEI(NetworkState networkState) {
		HashMap<NetworkState, Double> pX0X1 = p_X0_X1.get(networkState);
		return PHICompUtils.computeEI(pX0X1.values(), numOfNodes);
	}

	private final HashMap<NetworkState, Double> computeP_X0_X1(
			NetworkState networkState) {
		HashMap<NetworkState, Double> pX0X1 = new HashMap<NetworkState, Double>();

		List<NetworkState> states = SimNetUtils
				.getAllCertainNetworkStates(numOfNodes);

		for (NetworkState state : states) {
			NetworkState nextState = network.nextNetworkState(new NetworkState(
					state));
			double prob2u1 = PHICompUtils.getProbability2CertainState(nextState
					.getNodeStates(), networkState.getNodeStates(), totalPart);
			if (prob2u1 != 0)
				pX0X1.put(state, prob2u1);
		}

		return pX0X1;
	}

	private final HashMap<NetworkState, Double> computeP_M0_M1(
			NetworkState u1, List<Integer> part) {
		HashMap<NetworkState, Double> pM0M1 = new HashMap<NetworkState, Double>();
		int numOfNodesInPart = part.size();
		List<NodeState[]> states = SimNetUtils
				.getAllCertainStates(numOfNodesInPart);
		NodeState[] m0 = new NodeState[this.numOfNodes];
		NetworkState m0state = new NetworkState();
		
		for (int i = 0; i < numOfNodes; i++) {
			if (!part.contains(i))
				m0[i] = NodeState.MAXIMUM_ENTROPY;
		}
		for (NodeState[] state : states) {
			for (int i = 0; i < numOfNodesInPart; i++) {
				m0[part.get(i)] = state[i];
			}
			m0state.setNodeStates(m0);
			NetworkState nextState = network.nextNetworkState(m0state, part);
			double prob2u1 = PHICompUtils.getProbability2CertainState(nextState
					.getNodeStates(), u1.getNodeStates(), part);
			if (prob2u1 != 0)
				pM0M1.put(new NetworkState(m0state), prob2u1);
		}
//		System.out.println(pM0M1);
		return pM0M1;
	}

	private final double computeEI(HashMap<NetworkState, Double> pX0X1,
			List<HashMap<NetworkState, Double>> pM0M1s) {
		if (pX0X1.size()==0)
			return Double.POSITIVE_INFINITY;
			
		HashMap<NetworkState, Double> pM0M1 = new HashMap<NetworkState, Double>();
		NodeState[] nodeStates = new NodeState[numOfNodes];
		for(int i=0;i<numOfNodes;i++)
			nodeStates[i] = NodeState.CERTAIN_SILENT;
		pM0M1.put(new NetworkState(nodeStates), 1.0);

		for (HashMap<NetworkState, Double> pM0M1X : pM0M1s) {
			Set<Map.Entry<NetworkState, Double>> entries = pM0M1.entrySet();
			pM0M1 = new HashMap<NetworkState, Double>();
			for (Map.Entry<NetworkState, Double> entry : entries) {
				for (Map.Entry<NetworkState, Double> entry2 : pM0M1X.entrySet()) {
					pM0M1.put(NetworkState.combine(entry.getKey(), entry2
							.getKey()), entry.getValue() * entry2.getValue());
				}
			}

		}
		
		return PHICompUtils.computeKLDiv(pX0X1,pM0M1);
	}

	private final int normalization(Partition<?> partition) {
		if (partition.numOfParts() == 1) {
			return partition.part(0).size();
		}
		int sizeOfSmallestPart = partition.part(0).size();
		for (int i = 1; i < partition.numOfParts(); i++) {
			if (partition.part(i).size() < sizeOfSmallestPart)
				sizeOfSmallestPart = partition.part(i).size();
		}
		return (partition.numOfParts() - 1) * sizeOfSmallestPart;
	}

}
